Approved with small suggestions.

I still prefer a list of axes as the public API for join_dims over start_axis, n_axes, but I'm not willing to have a fight about it.

Counterargument is that for the internal uses we do have it was more awkward to do it like this and and you can't really do non-consecutive axis so there's also no benefit.

So we were asking users to transform their range into a list of axis, just so we can go and undo that because we want the range anyway

Yeah you are right, that's why I don't want to fight.

But I think from a user perspective, it's more natural to pass the list. I'm remembering how confused I was by the pt.split API to form this opinion.

What about split? Seems completely different thing

It requires two arguments, splits_size and n_splits, when I would only have expected one. Similarly, I would only expect one argument for join_dims -- the dims I want to join.

But I looked it up to type this message and I see that n_splits is actually optional. But it also doesn't have a docstring, so that probably didn't help me back then.

n_splits was needed before static shape, as tensor_variable([1, 2, 3]) wasn't known to have length 3 without introspecting it.

It could have been avoided by expecting a sequence of scalars but they didn't go with that.

For join_dims I think list of axis is fine if we support arbitrary axis.

I strongly believe that if we ask for sequence of axis the user will expect they can be non-consecutive. So the API immediately forces the user to think about the valid cases. It's a range. We could rename it to join_dims_range?

I'm still partial to a more flexible join_dims where we do the transpose to align the the dims to be joined for the user. But note that wasn't needed for our purposes of pack.

Although then you need to decide where they end up, as there's no longer a reasonable default). xarray puts it at the end with stack but positions don't matter as much for them.

Ah forgot to say why I did it though. You can't obtain the expand_dims edge case with the axis argument. There's no equivalent for join_dims(x, start_axis=2, n_axes=0)

Unless you offer the output_axis argument that decides where the joined axis goes

I think this is fine as-is. We can revisit it if we really want a more complex operation. I think re-arranging the dimensions is out of scope for this Op. We're slowly just building back up to dimshuffle.

I think re-arranging the dimensions is out of scope for this Op. We're slowly just building back up to dimshuffle.

It would be done by the helper, not the Op, So still combining DimShuffle and JoinDims, just doing the boilerplate work for the user. API would be something like: #1844

Agree with leaving it to another PR, but we have to decide on it soon, as it will be breaking change (this one already is)